import cv2
import os
import random
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

# 读入ORL人脸图像库中所有图像
data_dir = "F:/QAQ/The ORL Database of Faces/data/ORL/att_faces"
face_images = []
for i in range(1, 41):
    for j in range(1, 11):
        image_path = os.path.join(data_dir, "s%d" % i, "%d.pgm" % j)
        face_images.append(cv2.imread(image_path, cv2.IMREAD_GRAYSCALE))

# 将每个人的前5张人脸作为训练集，后5张人脸作为测试集
train_images = face_images[0:200:5]
test_images = face_images[1:200:5]

# 将图像转换为一维向量
train_data = np.array([img.flatten() for img in train_images])
test_data = np.array([img.flatten() for img in test_images])

# 对不同的类别数 k 进行测试，计算识别率
for k in [3, 5, 7]:
    # 使用KMeans算法对训练集进行聚类
    kmeans = KMeans(n_clusters=k, n_init=10, random_state=42)
    kmeans.fit(train_data)

    # 计算200个测试人脸的正确识别率，及其中错误识别的人脸编号
    correct_count = 0
    error_list = []
    for i, test_image in enumerate(test_images):
        true_label = (i // 5) + 1 # 计算真实标签
        predicted_label = kmeans.predict([test_image.flatten()])[0] + 1
        if true_label == predicted_label:
            correct_count += 1
        else:
            error_list.append((true_label, i % 5 + 6, true_label, predicted_label))  # 修正错误列表中的参数
    accuracy = correct_count / len(test_images) * 100

    # 显示识别结果
    print("k=%d，正确识别率：%.2f%%" % (k, accuracy))

    # 显示错误识别的人脸图像及其预测结果
    for error in error_list:
        true_label, image_index, _, predicted_label = error
        test_image = test_images[image_index - 1]
        plt.imshow(test_image, cmap="gray")
        plt.title("True Label: %d, Predicted Label: %d" % (true_label, predicted_label))
        plt.axis("off")
        plt.show()

